#!/bin/bash

# OpenCloudOS 节点管理工具 - 压力测试脚本
# 目标: 验证多Agent并发处理能力

set -euo pipefail  # 严格模式

echo "🔥 OpenCloudOS 压力测试开始"
echo "==========================="
echo "📋 测试目标:"
echo "  - 并发处理能力验证 (80个并发命令)"
echo "  - 大规模Agent协同测试 (50个Agent模拟)"
echo "  - 延迟性能测试 (目标: <50ms)"
echo "  - 系统稳定性验证"
echo

# 清理函数
cleanup() {
    echo "🧹 清理测试环境..."
    # 清理进程
    pkill -f "node_server" 2>/dev/null || true
    pkill -f "node_agent" 2>/dev/null || true
    sleep 2
    pkill -9 -f "node_server" 2>/dev/null || true
    pkill -9 -f "node_agent" 2>/dev/null || true
    
    # 清理临时日志和结果文件 - 包含所有测试生成的文件
    rm -f stress_*.log stress_result_*.txt stress_agent_*.log performance_*.log
    rm -f stress_server.log stress_hostname_*.log stress_script_*.log
    rm -f stress_web_*.log stress_db_*.log
    rm -f nohup.out
    
    echo "✅ 环境清理完成"
}

# 错误处理
handle_error() {
    echo "❌ 测试在第 $1 行失败！"
    cleanup
    exit 1
}

trap 'handle_error $LINENO' ERR

# 清理之前的进程
echo "清理环境..."
cleanup
sleep 2

# 启动Server
echo "启动 Node Server..."
./target/release/node_server start > stress_server.log 2>&1 &
SERVER_PID=$!
sleep 3

echo "Server PID: $SERVER_PID"

# 启动多个Agent (模拟大规模集群节点)
echo "启动多个Agent..."
declare -a AGENT_PIDS=()

# 启动15个web组Agent
for i in {1..15}; do
    ./target/release/node_agent --groups web --token secure_token_123 > stress_agent_web_$i.log 2>&1 &
    AGENT_PIDS+=($!)
    echo "Web Agent $i PID: $!"
done

# 启动10个database组Agent  
for i in {1..10}; do
    ./target/release/node_agent --groups database --token secure_token_123 > stress_agent_db_$i.log 2>&1 &
    AGENT_PIDS+=($!)
    echo "Database Agent $i PID: $!"
done

# 启动10个cache组Agent
for i in {1..10}; do
    ./target/release/node_agent --groups cache --token secure_token_123 > stress_agent_cache_$i.log 2>&1 &
    AGENT_PIDS+=($!)
    echo "Cache Agent $i PID: $!"
done

# 启动10个混合组Agent
for i in {1..10}; do
    ./target/release/node_agent --groups web,cache,default --token secure_token_123 > stress_agent_mixed_$i.log 2>&1 &
    AGENT_PIDS+=($!)
    echo "Mixed Agent $i PID: $!"
done

# 启动5个默认组Agent
for i in {1..5}; do
    ./target/release/node_agent --groups default --token secure_token_123 > stress_agent_default_$i.log 2>&1 &
    AGENT_PIDS+=($!)
    echo "Default Agent $i PID: $!"
done

echo "总共启动了 ${#AGENT_PIDS[@]} 个Agent"
sleep 8

# 检查所有Agent是否正常运行
echo "检查Agent状态..."
running_count=0
for pid in "${AGENT_PIDS[@]}"; do
    if ps -p $pid > /dev/null; then
        running_count=$((running_count + 1))
    fi
done
echo "运行中的Agent: $running_count/${#AGENT_PIDS[@]}"

# 如果Agent启动失败过多，显示警告但继续测试
if [ $running_count -lt $((${#AGENT_PIDS[@]} / 2)) ]; then
    echo "⚠️  警告: 超过一半的Agent启动失败，可能影响测试结果"
    echo "   建议检查系统资源和网络连接"
fi

# 并发测试 - 批量发送不同类型的命令
echo
echo "🚀 开始并发测试..."
echo "=================="

# 记录开始时间
start_time=$(date +%s)

# 延迟测试 - 单独测试几个命令的响应时间
echo "🕒 延迟性能测试 (目标: <50ms)..."
declare -a latency_results=()
for i in {1..5}; do
    # 使用秒级时间戳，避免毫秒精度问题
    cmd_start=$(date +%s)
    timeout 10s ./target/release/node_server command get-hostname > /dev/null 2>&1
    cmd_result=$?
    cmd_end=$(date +%s)
    
    if [ $cmd_result -eq 0 ]; then
        latency=$((cmd_end - cmd_start))
        # 转换为毫秒（估算值）
        latency_ms=$((latency * 1000))
        latency_results+=($latency_ms)
        echo "  延迟测试 $i: ~${latency_ms}ms"
    else
        echo "  延迟测试 $i: 超时/失败"
    fi
done

# 计算平均延迟
if [ ${#latency_results[@]} -gt 0 ]; then
    total_latency=0
    for lat in "${latency_results[@]}"; do
        total_latency=$((total_latency + lat))
    done
    avg_latency=$((total_latency / ${#latency_results[@]}))
    echo "  平均延迟: ${avg_latency}ms (目标: <50ms)"
    if [ $avg_latency -lt 50 ]; then
        echo "  ✅ 延迟要求达标"
    else
        echo "  ⚠️  延迟超过目标值"
    fi
else
    echo "  ❌ 延迟测试全部失败"
    avg_latency=9999
fi

echo

# 批量发送hostname命令 (30个)
echo "发送30个hostname命令..."
declare -a bg_pids=()
for i in {1..30}; do
    (
        timeout 15s ./target/release/node_server command get-hostname > stress_hostname_$i.log 2>&1
        echo "hostname_$i: $?" > stress_result_hostname_$i.txt
    ) &
    bg_pids+=($!)
done

# 批量发送script命令 (20个) 
echo "发送20个script命令..."
for i in {1..20}; do
    (
        timeout 20s ./target/release/node_server command run-script \
            --script "echo 'Stress test #$i'; date; sleep 1; whoami" \
            --timeout 10 > stress_script_$i.log 2>&1
        echo "script_$i: $?" > stress_result_script_$i.txt
    ) &
    bg_pids+=($!)
done

# 批量发送分组命令 (30个)
echo "发送30个分组命令..."
for i in {1..10}; do
    (
        timeout 15s ./target/release/node_server command --group web get-hostname > stress_web_$i.log 2>&1
        echo "web_$i: $?" > stress_result_web_$i.txt
    ) &
    bg_pids+=($!)
done

for i in {1..10}; do
    (
        timeout 15s ./target/release/node_server command --group database run-script \
            --script "echo 'DB stress #$i'; uptime" --timeout 5 > stress_db_$i.log 2>&1
        echo "db_$i: $?" > stress_result_db_$i.txt
    ) &
    bg_pids+=($!)
done

for i in {1..10}; do
    (
        timeout 15s ./target/release/node_server command --group cache run-script \
            --script "echo 'Cache stress #$i'; free -m" --timeout 5 > stress_cache_$i.log 2>&1
        echo "cache_$i: $?" > stress_result_cache_$i.txt
    ) &
    bg_pids+=($!)
done

# 等待所有命令完成，最多等待60秒
echo "等待所有命令完成..."
wait_count=0
max_wait=60

while [ $wait_count -lt $max_wait ]; do
    all_done=true
    for pid in "${bg_pids[@]}"; do
        if kill -0 "$pid" 2>/dev/null; then
            all_done=false
            break
        fi
    done
    
    if [ "$all_done" = true ]; then
        echo "所有命令已完成"
        break
    fi
    
    sleep 1
    wait_count=$((wait_count + 1))
    
    # 每10秒显示进度
    if [ $((wait_count % 10)) -eq 0 ]; then
        echo "等待中... ($wait_count/$max_wait 秒)"
    fi
done

# 强制终止任何剩余的进程
for pid in "${bg_pids[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then
        echo "强制终止超时进程 $pid"
        kill -9 "$pid" 2>/dev/null || true
    fi
done

# 记录结束时间
end_time=$(date +%s)
duration=$((end_time - start_time))

echo
echo "📊 压力测试结果"
echo "==============="
echo "测试持续时间: ${duration}秒"

# 延迟性能总结
echo
echo "🕒 延迟性能总结:"
if [ ${#latency_results[@]} -gt 0 ]; then
    echo "  测试样本: ${#latency_results[@]} 个"
    echo "  平均延迟: ${avg_latency}ms"
    echo "  目标延迟: <50ms"
    if [ $avg_latency -lt 50 ]; then
        echo "  状态: ✅ 符合要求"
    else
        echo "  状态: ⚠️ 需要优化"
    fi
else
    echo "  状态: ❌ 测试失败"
fi

# 统计成功/失败的命令
success_count=0
fail_count=0

# 统计hostname命令结果
for i in {1..30}; do
    if [ -f "stress_result_hostname_$i.txt" ]; then
        result=$(cat "stress_result_hostname_$i.txt" | cut -d: -f2 | tr -d ' ')
        if [ "$result" = "0" ]; then
            success_count=$((success_count + 1))
        else
            fail_count=$((fail_count + 1))
        fi
    else
        fail_count=$((fail_count + 1))
    fi
done

# 统计script命令结果
for i in {1..20}; do
    if [ -f "stress_result_script_$i.txt" ]; then
        result=$(cat "stress_result_script_$i.txt" | cut -d: -f2 | tr -d ' ')
        if [ "$result" = "0" ]; then
            success_count=$((success_count + 1))
        else
            fail_count=$((fail_count + 1))
        fi
    else
        fail_count=$((fail_count + 1))
    fi
done

# 统计分组命令结果
for i in {1..10}; do
    if [ -f "stress_result_web_$i.txt" ]; then
        result=$(cat "stress_result_web_$i.txt" | cut -d: -f2 | tr -d ' ')
        if [ "$result" = "0" ]; then
            success_count=$((success_count + 1))
        else
            fail_count=$((fail_count + 1))
        fi
    else
        fail_count=$((fail_count + 1))
    fi
done

for i in {1..10}; do
    if [ -f "stress_result_db_$i.txt" ]; then
        result=$(cat "stress_result_db_$i.txt" | cut -d: -f2 | tr -d ' ')
        if [ "$result" = "0" ]; then
            success_count=$((success_count + 1))
        else
            fail_count=$((fail_count + 1))
        fi
    else
        fail_count=$((fail_count + 1))
    fi
done

for i in {1..10}; do
    if [ -f "stress_result_cache_$i.txt" ]; then
        result=$(cat "stress_result_cache_$i.txt" | cut -d: -f2 | tr -d ' ')
        if [ "$result" = "0" ]; then
            success_count=$((success_count + 1))
        else
            fail_count=$((fail_count + 1))
        fi
    else
        fail_count=$((fail_count + 1))
    fi
done

total_commands=$((success_count + fail_count))
success_rate=$(echo "scale=2; $success_count * 100 / $total_commands" | bc -l 2>/dev/null || echo "N/A")

echo "总命令数: $total_commands"
echo "成功: $success_count"
echo "失败: $fail_count"
echo "成功率: $success_rate%"

# 检查数据库中的任务结果
echo
echo "📊 数据库验证"
echo "============="
sleep 5  # 等待所有结果入库

# 简化数据库查询，避免复杂的Docker命令问题
cd deploy
if docker compose ps -q postgresql >/dev/null 2>&1; then
    db_results=$(docker exec -e PGPASSWORD=admin123 \
        $(docker compose ps -q postgresql) \
        psql -h localhost -U admin -d node_management \
        -t -c "SELECT COUNT(*) FROM task_results WHERE completed_at > (NOW() - INTERVAL '5 minutes');" 2>/dev/null | tr -d ' ' || echo "0")
    echo "数据库中近5分钟的任务结果: $db_results 条"
else
    echo "数据库查询跳过 (容器不可用)"
fi
cd ..

# 检查Agent状态
echo
echo "🤖 Agent状态检查"
echo "================"
final_running_count=0
for pid in "${AGENT_PIDS[@]}"; do
    if ps -p $pid > /dev/null; then
        final_running_count=$((final_running_count + 1))
    fi
done
echo "测试后仍运行的Agent: $final_running_count/${#AGENT_PIDS[@]}"

# 系统资源使用情况
echo
echo "💻 系统资源使用"
echo "==============="
echo "内存使用情况:"
ps aux | grep -E "(node_server|node_agent)" | grep -v grep | awk '{sum+=$4} END {printf "总内存使用: %.2f%%\n", sum}' || echo "无法获取内存信息"

echo "进程数量:"
echo "node_server: $(pgrep -c node_server || echo 0)"
echo "node_agent: $(pgrep -c node_agent || echo 0)"

# 清理
echo
echo "🧹 清理测试环境..."
kill $SERVER_PID 2>/dev/null || true
for pid in "${AGENT_PIDS[@]}"; do
    kill $pid 2>/dev/null || true
done
sleep 2

# 清理临时文件
rm -f stress_*.log stress_result_*.txt 2>/dev/null || true

echo
echo "🎉 压力测试完成！"
echo "================="
echo "📈 性能评估总结:"
echo "  - 并发测试: $total_commands 个命令"
echo "  - 成功率: $success_rate%"
echo "  - 平均延迟: ${avg_latency}ms (目标: <50ms)"
echo "  - Agent规模: ${#AGENT_PIDS[@]} 个 (模拟小规模集群)"
echo

if [ "$success_rate" != "N/A" ] && [ "$(echo "$success_rate >= 90" | bc -l 2>/dev/null || echo 0)" = "1" ] && [ $avg_latency -lt 50 ]; then
    echo "✅ 测试结果: 优秀 (成功率 >= 90%, 延迟 < 50ms)"
    exit 0
elif [ "$success_rate" != "N/A" ] && [ "$(echo "$success_rate >= 70" | bc -l 2>/dev/null || echo 0)" = "1" ] && [ $avg_latency -lt 100 ]; then
    echo "🟡 测试结果: 良好 (成功率 >= 70%, 延迟可接受)"
    exit 0
else
    echo "❌ 测试结果: 需要优化 (成功率过低或延迟过高)"
    exit 1
fi
